{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a78cafa5",
   "metadata": {},
   "source": [
    "# conv2d_transpose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "07566090",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "67062780",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conv2DTransposeWorkload(batch=1, height=4, width=4, in_filter=1024, out_filter=512, hkernel=4, wkernel=4, hpad=1, wpad=1, hstride=2, wstride=2, o_hpad=0, o_wpad=0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/media/pc/data/board/arria10/lxw/tasks/tvm-new/python/tvm/driver/build_module.py:280: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.\n",
      "  warnings.warn(\n",
      "2025-07-27 20:57:07.773 INFO load_module /tmp/tmpezcoxjdq/conv2d_transpose.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.0130457 sec/op, 82.3059 GOPS\n",
      "Conv2DTransposeWorkload(batch=1, height=8, width=8, in_filter=512, out_filter=256, hkernel=4, wkernel=4, hpad=1, wpad=1, hstride=2, wstride=2, o_hpad=0, o_wpad=0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 20:57:11.402 INFO load_module /tmp/tmpezcoxjdq/conv2d_transpose.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.0140443 sec/op, 76.4538 GOPS\n",
      "Conv2DTransposeWorkload(batch=1, height=16, width=16, in_filter=256, out_filter=128, hkernel=4, wkernel=4, hpad=1, wpad=1, hstride=2, wstride=2, o_hpad=0, o_wpad=0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 20:57:14.657 INFO load_module /tmp/tmpezcoxjdq/conv2d_transpose.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.0107416 sec/op, 99.9614 GOPS\n",
      "Conv2DTransposeWorkload(batch=1, height=4, width=4, in_filter=1024, out_filter=512, hkernel=4, wkernel=4, hpad=1, wpad=1, hstride=2, wstride=2, o_hpad=0, o_wpad=0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[20:57:19] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 20:57:20.194 INFO load_module /tmp/tmp7jf_ai96/conv2d_transpose.o\n",
      "[20:57:20] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/runtime/profiling.cc:101: Warning: No timer implementation for ext_dev, using default timer instead. It may be inaccurate or have extra overhead.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.152596 sec/op, 7.03651 GOPS\n",
      "Conv2DTransposeWorkload(batch=1, height=8, width=8, in_filter=512, out_filter=256, hkernel=4, wkernel=4, hpad=1, wpad=1, hstride=2, wstride=2, o_hpad=0, o_wpad=0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[20:57:23] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 20:57:24.185 INFO load_module /tmp/tmp7jf_ai96/conv2d_transpose.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.151765 sec/op, 7.07501 GOPS\n",
      "Conv2DTransposeWorkload(batch=1, height=16, width=16, in_filter=256, out_filter=128, hkernel=4, wkernel=4, hpad=1, wpad=1, hstride=2, wstride=2, o_hpad=0, o_wpad=0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[20:57:27] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 20:57:27.686 INFO load_module /tmp/tmp7jf_ai96/conv2d_transpose.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.151822 sec/op, 7.07237 GOPS\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "import pytest\n",
    "import numpy as np\n",
    "from collections import namedtuple\n",
    "\n",
    "import tvm\n",
    "from tvm import te\n",
    "from tvm import relay\n",
    "from tvm import autotvm\n",
    "from tvm.contrib import utils\n",
    "from tvm.contrib.pickle_memoize import memoize\n",
    "from tvm import topi\n",
    "import tvm.topi.testing\n",
    "import vta\n",
    "from vta import program_fpga, reconfig_runtime\n",
    "import vta.testing\n",
    "from vta.testing import simulator\n",
    "\n",
    "\n",
    "Workload = namedtuple(\n",
    "    \"Conv2DTransposeWorkload\",\n",
    "    [\n",
    "        \"batch\",\n",
    "        \"height\",\n",
    "        \"width\",\n",
    "        \"in_filter\",\n",
    "        \"out_filter\",\n",
    "        \"hkernel\",\n",
    "        \"wkernel\",\n",
    "        \"hpad\",\n",
    "        \"wpad\",\n",
    "        \"hstride\",\n",
    "        \"wstride\",\n",
    "        \"o_hpad\",\n",
    "        \"o_wpad\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Get batch info from env\n",
    "env = vta.get_env()\n",
    "\n",
    "# DCGAN workloads\n",
    "dcgan_wklds = [\n",
    "    # dcgan\n",
    "    (\"DCGAN.CT1\", Workload(env.BATCH, 4, 4, 1024, 512, 4, 4, 1, 1, 2, 2, 0, 0)),\n",
    "    (\"DCGAN.CT2\", Workload(env.BATCH, 8, 8, 512, 256, 4, 4, 1, 1, 2, 2, 0, 0)),\n",
    "    (\"DCGAN.CT3\", Workload(env.BATCH, 16, 16, 256, 128, 4, 4, 1, 1, 2, 2, 0, 0)),\n",
    "]\n",
    "\n",
    "# FIXME: we need a custom clip operator to circumvent a pattern detection limitation\n",
    "@tvm.te.tag_scope(tag=topi.tag.ELEMWISE)\n",
    "def my_clip(x, a_min, a_max):\n",
    "    \"\"\"Unlike topi's current clip, put min and max into two stages.\"\"\"\n",
    "    const_min = tvm.tir.const(a_min, x.dtype)\n",
    "    const_max = tvm.tir.const(a_max, x.dtype)\n",
    "    x = te.compute(x.shape, lambda *i: tvm.te.min(x(*i), const_max), name=\"clipA\")\n",
    "    x = te.compute(x.shape, lambda *i: tvm.te.max(x(*i), const_min), name=\"clipB\")\n",
    "    return x\n",
    "\n",
    "\n",
    "# Helper function to get factors\n",
    "def _find_factors(n):\n",
    "    factors = []\n",
    "    for f in range(1, n + 1):\n",
    "        if n % f == 0:\n",
    "            factors.append(f)\n",
    "    return factors\n",
    "\n",
    "\n",
    "def run_conv2d_transpose(\n",
    "    env, remote, wl, target, check_correctness=True, print_ir=False, samples=4\n",
    "):\n",
    "\n",
    "    # Workload assertions\n",
    "    assert wl.hpad == wl.wpad\n",
    "\n",
    "    # Perform packing only if we are targeting the accelerator\n",
    "    if \"arm_cpu\" in target.keys:\n",
    "        data_pack = False\n",
    "        layout = \"NCHW\"\n",
    "        fcompute = topi.arm_cpu.conv2d_transpose_nchw\n",
    "        fschedule = topi.arm_cpu.schedule_conv2d_transpose_nchw\n",
    "    elif \"vta\" in target.keys:\n",
    "        data_pack = True\n",
    "        layout = \"NCHW%dn%dc\" % (env.BATCH, env.BLOCK_IN)\n",
    "        fcompute = vta.top.conv2d_transpose_packed\n",
    "        fschedule = vta.top.schedule_conv2d_transpose_packed\n",
    "\n",
    "    # Derive shapes depending upon packing\n",
    "\n",
    "    a_shape = (wl.batch, wl.in_filter, wl.height, wl.width)\n",
    "    w_shape = (wl.in_filter, wl.out_filter, wl.hkernel, wl.wkernel)\n",
    "    if data_pack:\n",
    "        data_shape = (\n",
    "            wl.batch // env.BATCH,\n",
    "            wl.in_filter // env.BLOCK_IN,\n",
    "            wl.height,\n",
    "            wl.width,\n",
    "            env.BATCH,\n",
    "            env.BLOCK_IN,\n",
    "        )\n",
    "        kernel_shape = (\n",
    "            wl.out_filter // env.BLOCK_OUT,\n",
    "            wl.in_filter // env.BLOCK_IN,\n",
    "            wl.hkernel,\n",
    "            wl.wkernel,\n",
    "            env.BLOCK_OUT,\n",
    "            env.BLOCK_IN,\n",
    "        )\n",
    "    else:\n",
    "        data_shape = a_shape\n",
    "        kernel_shape = w_shape\n",
    "    data = te.placeholder(data_shape, name=\"data\", dtype=env.inp_dtype)\n",
    "    kernel = te.placeholder(kernel_shape, name=\"kernel\", dtype=env.wgt_dtype)\n",
    "    padding = relay.nn.get_pad_tuple2d((wl.hpad, wl.wpad))\n",
    "\n",
    "    # Define base computation schedule\n",
    "    with target:\n",
    "\n",
    "        res = fcompute(\n",
    "            data, kernel, (wl.hstride, wl.wstride), padding, env.acc_dtype, (wl.o_hpad, wl.o_wpad)\n",
    "        )\n",
    "        res = topi.right_shift(res, env.WGT_WIDTH)\n",
    "        res = my_clip(res, 0, (1 << env.OUT_WIDTH - 1) - 1)\n",
    "        res = topi.cast(res, env.out_dtype)\n",
    "        # Derive base schedule\n",
    "        s = fschedule([res])\n",
    "        if print_ir:\n",
    "            print(vta.lower(s, [data, kernel, res], simple_mode=True))\n",
    "\n",
    "    # Derive number of ops\n",
    "    fout_height = (wl.height - 1) * wl.hstride - 2 * wl.hpad + wl.hkernel + wl.o_hpad\n",
    "    fout_width = (wl.width - 1) * wl.wstride - 2 * wl.wpad + wl.wkernel + wl.o_wpad\n",
    "    num_ops = (\n",
    "        2\n",
    "        * wl.batch\n",
    "        * fout_height\n",
    "        * fout_width\n",
    "        * wl.hkernel\n",
    "        * wl.wkernel\n",
    "        * wl.out_filter\n",
    "        * wl.in_filter\n",
    "    )\n",
    "\n",
    "    # @memoize(\"vta.tests.test_benchmark_topi.conv2d.verify_nhwc\")\n",
    "    def get_ref_data():\n",
    "        # derive min max for act and wgt types (max non inclusive)\n",
    "        a_min, a_max = 0 - (1 << (env.INP_WIDTH - 1)), (1 << (env.INP_WIDTH - 1))\n",
    "        w_min, w_max = 0 - (1 << (env.WGT_WIDTH - 1)), (1 << (env.WGT_WIDTH - 1))\n",
    "        a_np = np.random.randint(a_min, a_max, size=a_shape).astype(data.dtype)\n",
    "        w_np = np.random.randint(\n",
    "            w_min, w_max, size=(wl.in_filter, wl.out_filter, wl.hkernel, wl.wkernel)\n",
    "        ).astype(kernel.dtype)\n",
    "        r_np = tvm.topi.testing.conv2d_transpose_nchw_python(\n",
    "            a_np.astype(env.acc_dtype),\n",
    "            w_np.astype(env.acc_dtype),\n",
    "            (wl.hstride, wl.wstride),\n",
    "            wl.hpad,\n",
    "            (wl.o_hpad, wl.o_wpad),\n",
    "        ).astype(env.acc_dtype)\n",
    "        return a_np, w_np, r_np\n",
    "\n",
    "    # Data in original format\n",
    "    data_np, kernel_np, res_ref = get_ref_data()\n",
    "    if data_pack:\n",
    "        data_np = data_np.reshape(\n",
    "            wl.batch // env.BATCH,\n",
    "            env.BATCH,\n",
    "            wl.in_filter // env.BLOCK_IN,\n",
    "            env.BLOCK_IN,\n",
    "            wl.height,\n",
    "            wl.width,\n",
    "        ).transpose((0, 2, 4, 5, 1, 3))\n",
    "        kernel_np = kernel_np.reshape(\n",
    "            wl.in_filter // env.BLOCK_IN,\n",
    "            env.BLOCK_IN,\n",
    "            wl.out_filter // env.BLOCK_OUT,\n",
    "            env.BLOCK_OUT,\n",
    "            wl.hkernel,\n",
    "            wl.wkernel,\n",
    "        ).transpose((2, 0, 4, 5, 3, 1))\n",
    "        kernel_np = np.flip(kernel_np, 2)\n",
    "        kernel_np = np.flip(kernel_np, 3)\n",
    "\n",
    "    # Build\n",
    "    if \"vta\" in target.keys:\n",
    "        with vta.build_config(disabled_pass={\"tir.CommonSubexprElimTIR\"}):\n",
    "            mod = vta.build(\n",
    "                s,\n",
    "                [data, kernel, res],\n",
    "                target=target,\n",
    "                target_host=env.target_host,\n",
    "                name=\"conv2d_transpose\",\n",
    "            )\n",
    "    else:\n",
    "        mod = tvm.build(\n",
    "            s,\n",
    "            [data, kernel, res],\n",
    "            target=target,\n",
    "            target_host=env.target_host,\n",
    "            name=\"conv2d_transpose\",\n",
    "        )\n",
    "    temp = utils.tempdir()\n",
    "    mod.save(temp.relpath(\"conv2d_transpose.o\"))\n",
    "    remote.upload(temp.relpath(\"conv2d_transpose.o\"))\n",
    "    f = remote.load_module(\"conv2d_transpose.o\")\n",
    "    dev = remote.device(str(target))\n",
    "\n",
    "    res_np = np.zeros(topi.utils.get_const_tuple(res.shape)).astype(res.dtype)\n",
    "    data_arr = tvm.nd.array(data_np, dev)\n",
    "    kernel_arr = tvm.nd.array(kernel_np, dev)\n",
    "    res_arr = tvm.nd.array(res_np, dev)\n",
    "    time_f = f.time_evaluator(\"conv2d_transpose\", dev, number=samples)\n",
    "\n",
    "    # In vta sim mode, collect simulator runtime statistics\n",
    "    stats = {}\n",
    "    cost = None\n",
    "    if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "        # Check if we're in local RPC mode (allows us to rebuild the\n",
    "        # runtime on the fly when varying the VTA designs)\n",
    "        local_rpc = int(os.environ.get(\"VTA_LOCAL_SIM_RPC\", \"0\"))\n",
    "        if local_rpc:\n",
    "            if env.TARGET == \"sim\":\n",
    "                remote.get_function(\"vta.simulator.profiler_clear\")()\n",
    "            else:\n",
    "                remote.get_function(\"vta.tsim.profiler_clear\")()\n",
    "            cost = time_f(data_arr, kernel_arr, res_arr)\n",
    "            if env.TARGET == \"sim\":\n",
    "                stats = json.loads(remote.get_function(\"vta.simulator.profiler_status\")())\n",
    "            else:\n",
    "                stats = json.loads(remote.get_function(\"vta.tsim.profiler_status\")())\n",
    "        else:\n",
    "            simulator.clear_stats()\n",
    "            cost = time_f(data_arr, kernel_arr, res_arr)\n",
    "            stats = simulator.stats()\n",
    "    else:\n",
    "        cost = time_f(data_arr, kernel_arr, res_arr)\n",
    "\n",
    "    # Check correctness\n",
    "    correct = False\n",
    "    if check_correctness:\n",
    "        res_orig = res_arr.numpy()\n",
    "        if data_pack:\n",
    "            res_orig = res_orig.transpose((0, 4, 1, 5, 2, 3)).reshape(\n",
    "                wl.batch, wl.out_filter, fout_height, fout_width\n",
    "            )\n",
    "        res_ref = res_ref >> env.WGT_WIDTH\n",
    "        res_ref = np.clip(res_ref, 0, (1 << env.OUT_WIDTH - 1) - 1)\n",
    "        res_ref = res_ref.astype(env.out_dtype)\n",
    "        correct = np.allclose(res_orig, res_ref)\n",
    "\n",
    "    gops = (num_ops / cost.mean) / float(10**9)\n",
    "    status = \"PASSED\" if correct else \"FAILED\"\n",
    "    if \"arm_cpu\" in target.keys:\n",
    "        device = \"CPU\"\n",
    "    elif \"vta\" in target.keys:\n",
    "        device = \"VTA\"\n",
    "    print(\"%s CONV2D TEST %s: Time cost = %g sec/op, %g GOPS\" % (device, status, cost.mean, gops))\n",
    "\n",
    "    return correct, cost, stats\n",
    "\n",
    "\n",
    "@pytest.mark.parametrize(\"device\", [\"vta\", \"arm_cpu\"])\n",
    "def test_conv2d_transpose(device):\n",
    "    def _run(env, remote):\n",
    "        if device == \"vta\":\n",
    "            target = env.target\n",
    "            if env.TARGET not in [\"sim\", \"tsim\"]:\n",
    "                assert tvm.runtime.enabled(\"rpc\")\n",
    "                program_fpga(remote, bitstream=None)\n",
    "                reconfig_runtime(remote)\n",
    "        elif device == \"arm_cpu\":\n",
    "            target = env.target_vta_cpu\n",
    "        with autotvm.tophub.context(target):  # load pre-tuned schedule parameters\n",
    "            for _, wl in dcgan_wklds:\n",
    "                print(wl)\n",
    "                run_conv2d_transpose(env, remote, wl, target)\n",
    "\n",
    "    vta.testing.run(_run)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_conv2d_transpose(device=\"arm_cpu\")\n",
    "    test_conv2d_transpose(device=\"vta\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
